# see https://github.com/pymupdf/PyMuPDF-Utilities/blob/master/examples/extract-images/extract-from-pages.py
import fitz
from io import BytesIO
from PIL import Image

PADDLEOCR_ARGS = {
    "show_log": False,  # 关闭 log
    "use_angle_cls": True,
    "use_gpu": False,  # 不使用 cpu
    "enable_mkldnn": True,  # 多核加速 https://aistudio.baidu.com/paddle/forum/topic/show/966851
}
DEBUG = True  # only affect main()


def recoverpix(doc, item):  # move to top for declaration
    xref = item[0]  # xref of PDF image5
    smask = item[1]  # xref of its /SMask

    # special case: /SMask or /Mask exists
    if smask > 0:
        pix0 = fitz.Pixmap(doc.extract_image(xref)["image"])
        if pix0.alpha:  # catch irregular situation
            pix0 = fitz.Pixmap(pix0, 0)  # remove alpha channel
        mask = fitz.Pixmap(doc.extract_image(smask)["image"])

        try:
            pix = fitz.Pixmap(pix0, mask)
        except:  # fallback to original base image in case of problems
            pix = fitz.Pixmap(doc.extract_image(xref)["image"])

        if pix0.n > 3:
            ext = "pam"
        else:
            ext = "png"

        return {  # create dictionary expected by caller
            "ext": ext,
            "image": pix.tobytes(ext),
        }

    # special case: /ColorSpace definition exists
    # to be sure, we convert these cases to RGB PNG images
    if "/ColorSpace" in doc.xref_object(xref, compressed=True):
        pix = fitz.Pixmap(doc, xref)
        pix = fitz.Pixmap(fitz.csRGB, pix)
        return {  # create dictionary expected by caller
            "ext": "png",
            "colorspace": 3,
            "image": pix.tobytes("png"),
        }
    return doc.extract_image(xref)


# 导出论文所有格式
def dumpFileFormat(doc: fitz.Document):
    # 导出论文标题
    def dumpFileTitle(doc: fitz.Document) -> str:
        if doc.metadata["title"]:
            return doc.metadata["title"]
        title = []

        fontMax = 0.0
        blocks = doc.load_page(0).get_text("dict")["blocks"]
        # 遍历每个块，寻找文字，跳过不存在文字的块，type=0就是图片，限定寻找的标题在半页内
        for block in filter(
            lambda block: block["type"] != 0, blocks[: int(len(blocks) * 0.5)]
        ):
            for line in block["lines"]:
                for span in line["spans"]:  # 迭代行内
                    if span["size"] > fontMax:
                        fontMax = span["size"]
                        title = [span["text"]]
                        continue
                    elif span["size"] == fontMax:
                        title.append(span["text"])
                    continue
        return "".join(title)

    def dumpFileAuthor(doc: fitz.Document) -> str:
        return "qwer"

    def dumpFileSubject(doc: fitz.Document) -> str:
        return "qwer"

    def dumpFileCreationDate(doc: fitz.Document):
        return "asdf"

    # This function exports the abstract section of a paper and returns it as a string.
    def dumpFileAbstract(doc: fitz.Document) -> str:
        abstractCN = []  # Stores Chinese abstract
        abstractEN = []  # Stores English abstract
        abstract = []  # Integrated abstract

        blockBox = []  # Stores block position information
        blockBox2 = []  # Stores sub-block position information
        # 0: Find block, 1: Collect Chinese block, 2: Collect English block, 3: Collection completed
        abstractFlag = 0

        # Traverse each block of the document, skip blocks without text (type=0 represents images)
        blocks = doc.load_page(0).get_text("dict")["blocks"]
        for block in filter(lambda block: block["type"] == 0, blocks):
            for line in block["lines"]:  # Iterating each line
                for span in line["spans"]:  # Iterating each span in the line
                    # Find abstract index block (Chinese or English)
                    temp = span["text"].replace(" ", "").lower()
                    if abstractFlag == 0:
                        if "摘要" in temp:
                            abstractFlag = 1
                            blockBox = span["bbox"]
                            continue
                        elif "abstract" in temp:
                            abstractFlag = 2
                            blockBox = span["bbox"]
                            continue

                    # Collect index block
                    if abstractFlag == 1:
                        if not blockBox2 and "】　" in span["text"]:
                            # Chinese text processing filters out brackets
                            continue
                        if (
                            not blockBox2
                            and abs(span["bbox"][0] - blockBox[0]) < 40
                            or blockBox2
                            and blockBox[0] - 40 < span["bbox"][0] < blockBox2[2] + 40
                            and abs(span["bbox"][1] - blockBox2[1]) < 40
                        ):
                            blockBox2 = span["bbox"]
                            if "　" in span["text"]:
                                span["text"] = "："
                            abstractCN.append(span)
                        elif blockBox2:
                            blockBox2 = []
                            abstractFlag = 0

                    if abstractFlag == 2:
                        if (
                            not blockBox2
                            and abs(span["bbox"][0] - blockBox[0]) < 40
                            or blockBox2
                            and blockBox[0] - 40 < span["bbox"][0] < blockBox2[2] + 40
                            and abs(span["bbox"][1] - blockBox2[1]) < 40
                        ):
                            blockBox2 = span["bbox"]
                            abstractEN.append(span)

        # Format collected abstract span
        if abstractCN:
            blockBox = abstractCN[0]["bbox"]
            abstract.append(abstractCN[0]["text"])
            for span in abstractCN[1:]:
                if abs(span["bbox"][1] - blockBox[1]) < 5:
                    # If the y value of a span is within +/-5 pixels of the previous span,
                    # it is considered on the same line and added directly to abstract
                    abstract.append(span["text"])
                elif abs(span["bbox"][1] - blockBox[1]) < 20:
                    blockBox = span["bbox"]
                    if "【" in span["text"]:
                        # If it encounters the next type of block, ending is skipped
                        break
                    if abstract[-1][-1] != "-":
                        # If the previous and next line only contain letters, a space is added before appending
                        abstract.append(" ")
                    abstract.append(span["text"])
                else:
                    break

        # Format collected abstract span
        if abstractEN:
            blockBox = abstractEN[0]["bbox"]
            abstract.append(abstractEN[0]["text"])
            for span in abstractEN[1:]:
                if abs(span["bbox"][1] - blockBox[1]) < 5:
                    # If the y value of a span is within +/-5 pixels of the previous span,
                    # it is considered on the same line and added directly to abstract
                    abstract.append(span["text"])
                elif abs(span["bbox"][1] - blockBox[1]) < 12:
                    blockBox = span["bbox"]
                    if abstract[-1][-1] != "-":
                        # If the previous line only contains letters and the next line begins with letters, a space is added before appending
                        abstract.append(" ")
                    abstract.append(span["text"])
                else:
                    break
        return "".join(abstract)

    # 导出论文关键词
    def dumpFileKeywords(doc: fitz.Document):
        keywordsCN = []
        keywordsEN = []
        keywords = []
        blockBox = []
        blockBox2 = []
        # 0:找块，1:收集中文块，2:收集英文块，3:收集完成
        keywordsFlag = 0

        # 遍历每个块，寻找文字，跳过不存在文字的块，type=0就是图片
        blocks = doc.load_page(0).get_text("dict")["blocks"]
        for block in filter(lambda block: block["type"] == 0, blocks):
            for line in block["lines"]:  # 迭代行
                for span in line["spans"]:  # 迭代行内
                    # 寻找摘要索引块（中、英文）
                    temp = span["text"].replace(" ", "").lower()
                    if keywordsFlag == 0:
                        if "关键词" in temp:
                            keywordsFlag = 1
                            blockBox = span["bbox"]
                            continue
                        elif "keywords" in temp:
                            keywordsFlag = 2
                            blockBox = span["bbox"]
                            continue

                    # 收集索引块
                    if keywordsFlag == 1:
                        if not blockBox2 and "】　" in span["text"]:
                            # 当开头时，中文处理过滤方括号
                            continue
                        if (  # span['bbox'][0] == blockBox[0] or  # 与开头在同一列上的段
                            # 当开头时，与索引块在垂直距离小于20、水平距离小于50的段（一行间、两行间）
                            not blockBox2
                            and (
                                blockBox[0] - 50 < span["bbox"][0]
                                and span["bbox"][0] < blockBox[2] + 50
                                and abs(span["bbox"][1] - blockBox[1]) < 20
                            )
                            or blockBox2
                            and (
                                blockBox[0] - 40 < span["bbox"][0]
                                and span["bbox"][0] < blockBox2[2] + 40
                                and abs(span["bbox"][1] - blockBox2[1]) < 20
                            )
                        ):  # 非开头时，与上在垂直距离小于20、水平距离小于40的段（一行间、两行间）
                            blockBox2 = span["bbox"]
                            keywordsCN.append(span)
                        elif blockBox2:
                            # 如果开过头了，没发现范围内文本，完成一次收集
                            blockBox2 = []
                            keywordsFlag = 0

                    if keywordsFlag == 2:
                        if not blockBox2 and "】　" in span["text"]:
                            # 当开头时，中文处理过滤方括号
                            continue
                        if (  # span['bbox'][0] == blockBox[0] or  # 与开头在同一列上的段
                            # 当开头时，与索引块在垂直距离小于20、水平距离小于50的段（一行间、两行间）
                            not blockBox2
                            and (
                                blockBox[0] - 50 < span["bbox"][0]
                                and span["bbox"][0] < blockBox[2] + 50
                                and abs(span["bbox"][1] - blockBox[1]) < 20
                            )
                            or blockBox2
                            and (
                                blockBox[0] - 40 < span["bbox"][0]
                                and span["bbox"][0] < blockBox2[2] + 40
                                and abs(span["bbox"][1] - blockBox2[1]) < 20
                            )
                        ):  # 非开头时，与上在垂直距离小于20、水平距离小于40的段（一行间、两行间）
                            blockBox2 = span["bbox"]
                            keywordsEN.append(span)
                        elif blockBox2:
                            # 如果开过头了，没发现范围内文本，完成一次收集
                            blockBox2 = []
                            keywordsFlag = 0

        # 整理收集的摘要span
        if keywordsCN:
            blockBox = keywordsCN[0]["bbox"]
            keywords.append(keywordsCN[0]["text"])
            for span in keywordsCN[1:]:
                if abs(span["bbox"][1] - blockBox[1]) < 5:
                    # 一行内span y在+-5个像素以内可以认为是在同一行上，加入到keywords
                    keywords.append(span["text"])
                elif abs(span["bbox"][1] - blockBox[1]) < 20:
                    # 不是一行的就换行，一行内span y在+-20个像素以内可以认为是在同一段上
                    blockBox = span["bbox"]
                    if "【" in span["text"]:
                        # 遇到下一类块，省去结束
                        break
                    if keywords[-1][-1] != "-":
                        # 上行行末和下行行首都是字母的，加一个空格
                        keywords.append(" ")
                    keywords.append(span["text"])
                else:
                    # 不是同一段的就省去
                    break

        # 整理收集的摘要span
        if keywordsEN:
            blockBox = keywordsEN[0]["bbox"]
            keywords.append(keywordsEN[0]["text"])
            for span in keywordsEN[1:]:
                if abs(span["bbox"][1] - blockBox[1]) < 5:
                    # 一行内span y在+-5个像素以内可以认为是在同一行上，加入到keywords
                    keywords.append(span["text"])
                elif abs(span["bbox"][1] - blockBox[1]) < 12:
                    # 不是一行的就换行，一行内span y在+-12个像素以内可以认为是在同一段上
                    blockBox = span["bbox"]
                    if keywords[-1][-1] != "-":
                        # 上行行末和下行行首都是字母的，加一个空格
                        keywords.append(" ")
                    keywords.append(span["text"])
                else:
                    # 不是同一段的就省去
                    break

        # 去除空白项
        return " ".join(filter(lambda a: a != " ", keywords)).replace(";　", "")

    meta = {
        "title": dumpFileTitle(doc),
        "author": dumpFileAuthor(doc),
        "subject": dumpFileSubject(doc),
        "keywords": dumpFileKeywords(doc),
        "creationDate": dumpFileCreationDate(doc),
        "abstract": dumpFileAbstract(doc),
    }
    return meta


def ocr_zh(pil_img: Image.Image, xy=None, lang="ch"):  # xy: prefix
    from paddleocr import PaddleOCR
    import cv2
    import numpy as np

    ocr = PaddleOCR(lang=lang, **PADDLEOCR_ARGS)  # 关闭 debug output
    img = np.array(pil_img.convert("RGB"))
    result = ocr.ocr(img)[0]
    charpos = []
    for char in result:
        x, y = char[0][0]
        w = char[0][1][0] - char[0][0][0]
        h = char[0][3][1] - char[0][0][0]
        if xy is not None:
            x += xy[0]
            y += xy[1]
        charpos.append((char[1][0], x, y, w, h))
    return charpos


def main():
    """
    handle file entry from CLI.
    would import os, sys.
    and produce
    {filename}-out.txt
    {filename}-out-img%03d.png
    {filename}-out-img%03d.txt
    at where source file placed.
    To use this module, import pdftotext() instead.

    """
    import os, sys
    from json import dump

    doc = sys.argv[1]
    docname = doc[-4]
    data = pdftotext(open(doc, "rb").read())
    for num, x in enumerate(data[0], 1):  # pages
        if DEBUG:
            print(f"writing pagetext {num}")
        with open(f"{docname}.raw{num:03}.txt", "w") as f:
            f.write(x)
    for img in data[1]:  # imgs
        ref03 = f"{img[1][0]:03}"
        if DEBUG:
            print(f"writing img {ref03}")
        with open(f"{docname}.img{ref03}.{img[1][1]}", "wb") as out:
            out.write(img[0])
        with open(f"{docname}.img{ref03}.data.json", "w") as out:
            dump({"num": f"{ref03}", "data": img[2]}, out)


def pdftotext(pdf_bytes: BytesIO):
    """
    return:
    (
        pages: text per page in list,
        imgs: [
            (
                img: bytes,
                (xrefnum: int, extension: str),
                boxes: (
                    text: str,
                    x, y, w, h: number
                )
            ), ...
        ]
    )
    the length of image_list equals to image number
    there is no dict in return data
    """
    doc = fitz.open(stream=pdf_bytes)  # open document
    print("Dumping file format...")
    docformat = dumpFileFormat(doc)
    pgs = []  # pdf plain text each page
    imgs = []
    ischi = None
    for pagenum, page in enumerate(doc, 1):  # iterate the document pages
        out = ""
        print("dumping page", pagenum)  # TODO: add progress
        text = page.get_text()  # get plain text (is in UTF-8)
        out += text  # write text of page
        out += "\n"  # write page delimiter (form feed 0x0C)
        pgs.append(out)

        images = doc.get_page_images(pagenum - 1, full=True)
        imgxrefset = set()
        for imgref in images:
            # img: xref, smask, width, height
            if imgref[0] in imgxrefset:
                continue  # do not dump duplicated image
            img = recoverpix(doc, imgref)
            imgxrefset.add(imgref[0])
            image = img["image"]  # img in bytes
            pil_img = Image.open(BytesIO(img["image"]))
            # 4.12: coordinate convert: convert image coordination to abs coordination in pdf
            xy = page.get_image_bbox(  # https://pymupdf.readthedocs.io/en/latest/page.html#Page.get_image_bbox
                imgref
            )
            print(f"OCRing image {imgref[0]:03}")
            boxes = ocr_zh(pil_img, xy)  # ocr
            imgs.append(
                (image, (imgref[0], img["ext"]), boxes)
            )  # imgref for save in order,
    return (pgs, imgs)


if __name__ == "__main__":
    main()
